# ENTITY_LABELS = 
# ["O", "B-PER", "I-PER", "B-ORG", "I-ORG", "B-LOC", "I-LOC"]
# RELATION_LABELS = ["attribute", "belong", "contain", "unknown"]
# PROMPT_TEMPLATES = [
#    "[E1] is the [MASK] of [E2].",
#    "[E1]'s relation to [E2] is [MASK].",
#    "The relationship between [E1] and [E2] is [MASK]."
#]
SAMPLE = [
    {
        "text": "What is the age of the man?",
        "ner_labels": [0, 0, 0, 0, 0, 0, 0],  
        "entity_pair": ("age", "man"),
        "relation": "attribute"
    },
    {
        "text": "how many contracts does the person have?",
        "ner_labels": [0, 0, 1, 0, 0, 1, 0],  
        "entity_pair": ("contracts", "person"),
        "relation": "attribute"
    },
    {
        "text": "The major of Artificial Intelligence is one of the majors under the Department of Computer Science.",
        "ner_labels": [0, 0, 0, 3, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 4 ],  
        "entity_pair": ("Artificial Intelligence", "Department of Computer Science"),
        "relation": "belong"
    },
    {
        "text": "Query the quantity of contracts of the digitalization department",
        "ner_labels": [0, 0, 0, 0, 0, 0, 0, 3, 4],  
        "entity_pair": ("digitalization department","contracts"),
        "relation": "contain"
    },
    {
        "text": "Bob is an employee of the company.",
        "ner_labels": [1, 0, 0, 0, 0, 0, 0],  
        "entity_pair": ("company", "Bob"),
        "relation": "contain"
    },
    {
        "text": "What is the salary range of the engineer?",
        "ner_labels": [0, 0, 0, 0, 0, 0, 0, 0, 0],  
        "entity_pair": ("salary range", "engineer"),
        "relation": "attribute"
    },
    {
        "text": "How many projects does the marketing team manage?",
        "ner_labels": [0, 0, 0, 0, 0, 3, 4, 0],  
        "entity_pair": ("projects", "marketing team"),
        "relation": "attribute"
    },

    # 关系: belong
    {
        "text": "The program of Data Science belongs to the Faculty of Engineering.",
        "ner_labels": [0, 0, 0, 3, 4, 0, 0, 0, 3, 4, 4],  
        "entity_pair": ("Data Science", "Faculty of Engineering"),
        "relation": "belong"
    },
    {
        "text": "Machine Learning is a research field within the AI Lab.",
        "ner_labels": [3, 4, 0, 0, 0, 0, 0, 0, 3, 4],  
        "entity_pair": ("Machine Learning", "AI Lab"),
        "relation": "belong"
    },

    # 关系: contain
    {
        "text": "The finance division holds all budget reports.",
        "ner_labels": [0, 3, 4, 0, 0, 0, 0, 0],  
        "entity_pair": ("finance division", "budget reports"),
        "relation": "contain"
    },
    {
        "text": "Alice oversees the HR department's recruitment plans.",
        "ner_labels": [1, 0, 0, 3, 4, 0, 0],  
        "entity_pair": ("HR department", "recruitment plans"),
        "relation": "contain"
    },

    # 混合实体类型（LOC）
    {
        "text": "The Tokyo office manages the Asia-Pacific region.",
        "ner_labels": [0, 5, 6, 0, 0, 0, 5, 6],  
        "entity_pair": ("Tokyo office", "Asia-Pacific region"),
        "relation": "contain"
    },
    {
        "text": "What is the population of New York City?",
        "ner_labels": [0, 0, 0, 0, 5, 6, 0],  
        "entity_pair": ("population", "New York City"),
        "relation": "attribute"
    }
]